import { cloneDeep } from '@/utils'

interface DataItem {
  id: number
  children?: DataItem[]
  [key: string]: any
}

interface MapItem extends DataItem {
  pid?: number
  isLeftNode: boolean
  childrenSize: number
}

const setMap = (
  data: DataItem[],
  pid?: number,
  map = new Map<number, MapItem>()
) => {
  if (!Array.isArray(data)) return map

  data.forEach((item) => {
    map.set(item.id, {
      ...item,
      pid,
      isLeftNode: !item.children,
      childrenSize: (item.children ?? []).length
    })
    if (item.children) {
      setMap(item.children, item.id, map)
    }
  })

  return map
}

const getSelectedHandle = (
  ids: number[],
  map: Map<number, MapItem>,
  allIds: number[]
): any => {
  let pidList: number[] = []
  const selectedPids: number[] = []

  ids.forEach((id) => {
    const item = map.get(id)
    if (!item) return

    const pid = item.pid
    if (pid) pidList.push(pid)

    //去重
    pidList = [...new Set(pidList)]
    const parentItem = pid ? map.get(pid) : undefined
    if (parentItem) parentItem.childrenNum = (parentItem.childrenNum ?? 0) + 1
  })

  pidList.forEach((id) => {
    const item = map.get(id)
    if (!item) return

    if (item.childrenNum === item.childrenSize) {
      selectedPids.push(id)
    }
  })

  if (!selectedPids.length) return allIds

  allIds = allIds.concat(selectedPids)
  return getSelectedHandle(selectedPids, map, allIds)
}

export interface TreeHandler {
  getLeftNode: (ids: number[]) => number[]
  getHalfSelected: (ids: number[]) => number[]
  getSelected: (ids: number[]) => number[]
}

export const useTreeHandler = (data: DataItem[]): TreeHandler => {
  const map = setMap(data)

  // ids => leftNode.id
  const getLeftNode = (ids: number[]) =>
    ids.filter((v) => map.get(v)?.isLeftNode)

  // ids => ids + parent.id (父级半选)
  const getHalfSelected = (ids: number[]) => {
    const queue = [...ids]
    const _ids = new Set()

    while (queue.length) {
      const id = queue.shift()!
      const node = map.get(id)
      if (node) {
        _ids.add(node.id)

        if (node.pid) {
          queue.push(node.pid)
        }
      }
    }

    return [..._ids.values()] as number[]
  }

  // ids => ids + parent.id (父级选中)
  const getSelected = (ids: number[]) => {
    const list = getSelectedHandle(ids, cloneDeep(map), ids)
    return list
  }

  return {
    getLeftNode,
    getHalfSelected,
    getSelected
  }
}
